import os

from fastapi import FastAPI
from langchain_core.prompts.chat import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain_community.llms import Tongyi
from langserve import add_routes
from zhipu_generate_token import generate_token

app = FastAPI(
    title="Multilingual Translation Service",
    version="1.0",
    description="This service provides multilingual translation using three different models: ZhipuAI, OpenAI, and Tongyi.",
)

system_message_prompt = """You are a translation expert proficient in various languages, \
and your task is to translate the content input by the user into {target_language} and return the translated content. \
Don't do other tasks."""

chat_prompt_template = ChatPromptTemplate.from_messages([
    ('system', system_message_prompt),
    ('user', '{input}')
])

zhipuai_model = ChatOpenAI(
    model_name='glm-4',
    openai_api_base='https://open.bigmodel.cn/api/paas/v4',
    openai_api_key=generate_token(os.getenv('ZHIPUAI_API_KEY'), 24*60*60),  # 设置 token 的过期时间为 1 天
)

openai_model = ChatOpenAI()

tongyi_model = Tongyi()

add_routes(
    app,
    chat_prompt_template | openai_model,
    path="/TranslateByOpenAI",
)

add_routes(
    app,
    chat_prompt_template | zhipuai_model,
    path="/TranslateByZhipuAI",
)

add_routes(
    app,
    chat_prompt_template | tongyi_model,
    path="/TranslateByTongyi",
)

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="localhost", port=8000)